import numpy as np

"""
功能：用变分模态分解(VMD)分离信号成分
技术：VMD算法实现模态分解
比喻：把混合果汁分离成纯苹果汁、橙汁等成分
"""

from vmdpy import VMD  # 需安装vmdpy库


def separate_components(signal, alpha=3000, tau=0.3, K=5): # 增大alpha
    """
    参数说明：
    signal : 原始信号 -> 混合果汁
    alpha  : 数据保真度权重 -> 分离力度 (建议1000~5000)
    tau    : 时间步长 -> 离心机转速 (通常0.1~0.5)
    K      : 模态数量 -> 想分离出几种果汁

    返回：K个本征模态分量
    """
    # 执行VMD分解（启动离心机）
    u, omega, _ = VMD(
        signal,
        alpha=alpha,
        tau=tau,
        K=K,
        DC=0,  # 禁用直流分量
        init=1,  # 初始化方式：均匀分布
        tol=1e-7
    )
    # return u
    # 质量检测（检查果汁纯度）：确保每个成分都是纯净色
    for i in range(K):
        if np.max(u[i]) < 0.1 * np.max(signal):
            print(f"警告：模态{i}可能是噪声残留！建议调整K值")
    return u
